package com.googlecode.connectlet.vspkm;

import java.security.cert.X509Certificate;
import java.security.interfaces.RSAPublicKey;

import sun.security.pkcs.PKCS7;

import com.googlecode.connectlet.Filter;
import com.googlecode.connectlet.crypto.cipher.rc5.RC564Cipher;
import com.googlecode.connectlet.crypto.random.RC4PrngCipher;
import com.googlecode.connectlet.crypto.rsa.RSA;
import com.googlecode.connectlet.crypto.stream.CfbDecipheror;
import com.googlecode.connectlet.crypto.stream.CfbEncipheror;
import com.googlecode.connectlet.util.ByteArrayQueue;
import com.googlecode.connectlet.util.Bytes;
import com.googlecode.connectlet.util.secure.CertKey;
import com.googlecode.connectlet.util.secure.CertMap;

public class VspkmFilter extends Filter {
	private static RC4PrngCipher random = RC4PrngCipher.getInstance();

	private static final int STATUS_PKCS7_HEAD = 0;
	private static final int STATUS_PKCS7_BODY = 1;
	private static final int STATUS_KEY_IV = 2;
	private static final int STATUS_CONNECTED = 3;

	private CertKey certKey;
	private CertMap certMap;
	private CfbEncipheror outCipher;
	private CfbDecipheror inCipher;
	private X509Certificate[] peerCertificates = null;
	private int status = STATUS_PKCS7_HEAD;

	public VspkmFilter(CertKey certKey, CertMap certMap) {
		this.certKey = certKey;
		this.certMap = certMap;
	}

	private ByteArrayQueue baqToSend = new ByteArrayQueue();
	private ByteArrayQueue baqRecv = new ByteArrayQueue();
	private int bytesToRecv = 4;

	@Override
	protected void send(byte[] b, int off, int len) {
		if (status != STATUS_CONNECTED) {
			baqToSend.add(b, off, len);
			return;
		}
		super.send(outCipher.update(b, off, len), 0, len);
	}

	@Override
	protected void onRecv(byte[] b, int off, int len) {
		// Connected
		if (status == STATUS_CONNECTED) {
			super.onRecv(inCipher.update(b, off, len), 0, len);
			return;
		}

		// Handshake
		baqRecv.add(b, off, len);
		if (baqRecv.length() < bytesToRecv) {
			return;
		}

		// 1. Got PKCS7 Header (4 Bytes)
		if (status == STATUS_PKCS7_HEAD) {
			byte[] buffer = baqRecv.array();
			int offset = baqRecv.offset();
			if (buffer[offset] != 0x30) {
				getConnection().disconnect();
				return;
			}
			int lenData = buffer[offset + 1] & 0xff;
			int lenLen = 0;
			if (lenData < 2) {
				getConnection().disconnect();
				return;
			}
			if (lenData > 128) {
				lenLen = lenData - 128;
				if (lenLen > 2) {
					getConnection().disconnect();
					return;
				}
				lenData = 0;
				for (int i = 0; i < lenLen; i ++) {
					lenData *= 256;
					lenData += buffer[offset + 2 + i] & 0xff;
				}
			}

			bytesToRecv = 2 + lenLen + lenData;
			status = STATUS_PKCS7_BODY;
			if (baqRecv.length() < bytesToRecv) {
				return;
			}
		}

		// 2. Got PKCS7 Body (Length Calculated in STATUS_PKCS7_HEAD)
		if (status == STATUS_PKCS7_BODY) {
			// 2-1. Get Certificate Chain
			byte[] pkcs7Bytes = new byte[bytesToRecv];
			baqRecv.remove(pkcs7Bytes);
			PKCS7 pkcs7;
			try {
				pkcs7 = new PKCS7(pkcs7Bytes);
			} catch (Exception e) {
				getConnection().disconnect();
				return;
			}
			if (certMap == null) {
				peerCertificates = pkcs7.getCertificates();
			} else {
				peerCertificates = certMap.getCertificateChain(pkcs7.getCertificates());
			}
			// 2-2. Verify Certificate Chain
			if (peerCertificates == null || peerCertificates.length == 0) {
				getConnection().disconnect();
				return;
			}
			if (certMap != null && !certMap.verify(peerCertificates)) {
				getConnection().disconnect();
				return;
			}
			// 2-3. Generate RC5-CFB Cipher with Random Key and IV 
			byte[] key = new byte[32];
			random.output(key);
			byte[] iv = new byte[16];
			random.output(iv);
			outCipher = new CfbEncipheror(new RC564Cipher(key), iv);
			// 2-4. Send Encrypted Key & IV
			RSAPublicKey publicKey = (RSAPublicKey) peerCertificates[0].getPublicKey();
			byte[] keyIv = Bytes.add(RSA.encrypt(publicKey, key), iv);
			super.send(keyIv, 0, keyIv.length);

			bytesToRecv = RSA.getBlockSize(certKey.getKey()) + 16;
			status = STATUS_KEY_IV;
			if (baqRecv.length() < bytesToRecv) {
				return;
			}
		}

		// 3. Got Encrypted Key & IV
		int blockSize = bytesToRecv - 16;
		try {
			byte[] rsaBytes = new byte[blockSize];
			baqRecv.remove(rsaBytes);
			byte[] key = RSA.decrypt(certKey.getKey(), rsaBytes);
			byte[] ivBytes = new byte[16];
			baqRecv.remove(ivBytes);
			inCipher = new CfbDecipheror(new RC564Cipher(key), ivBytes);
		} catch (Exception e) {
			getConnection().disconnect();
			return;
		}

		status = STATUS_CONNECTED;
		super.onConnect();
		if (baqRecv.length() > 0) {
			onRecv(baqRecv.array(), baqRecv.offset(), baqRecv.length());
		}
		baqRecv = null; // Never Used when Connected
		if (baqToSend.length() > 0) {
			send(baqToSend.array(), baqToSend.offset(), baqToSend.length());
		}
		baqToSend = null; // Never Used when Connected
	}

	@Override
	protected void onConnect() {
		ByteArrayQueue baqPkcs7 = new ByteArrayQueue();
		try {
			CertMap.exportPkcs7(certKey.getCertificateChain()).
					encodeSignedData(baqPkcs7.getOutputStream());
		} catch (Exception e) {
			getConnection().disconnect();
			return;
		}
		super.send(baqPkcs7.array(), baqPkcs7.offset(), baqPkcs7.length());
	}

	public X509Certificate[] getPeerCertificates() {
		return peerCertificates;
	}
}